import warnings

from ultralytics import YOLO

warnings.filterwarnings('ignore')

if __name__ == '__main__':
    model = YOLO('./dataset/yolo11s-seg.pt')
    results = model.train(
        data='./dataset/data.yaml',  # 数据集配置文件的路径
        project='./result',
        epochs=200,                  # 训练轮次总数
        save_period=10,              # 每10次存一个模型
        batch=16,                    # 批量大小，即单次输入多少图片训练
        imgsz=640,                   # 训练图像尺寸
        workers=8,                  # 加载数据的工作线程数
        device=0,                    # 指定训练的计算设备，无nvidia显卡则改为 'cpu'
        patience=30,                 # 早停轮次（验证指标未提升则终止）
        weight_decay=0.0005,         # 权重衰减系数（防过拟合）
        label_smoothing=0.05,        # 标签平滑强度（建议6分类设0.05-0.1）
        optimizer='AdamW',           # 训练使用优化器，可选 auto,SGD,Adam,AdamW 等
        amp=True,                    # True 或者 False, 解释为：自动混合精度(AMP) 训练
        cos_lr=True,                 # 余弦退火学习率缓解过拟合‌
        cache=False                  # True 在内存中缓存数据集图像，服务器推荐开启
    )

